#include <iostream>

#include <vector>
#include <map>
#include <set>
#include <memory>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <lemon/list_graph.h>
#include <lemon/bin_heap.h>
#include <lemon/unionfind.h>
#include <lemon/dim2.h>

#include <cluster.h>

#if !defined(DIST_AVG) && !defined(DIST_DIV) && !defined(DIST_MAH) && \
  !defined(DIST_BAT) && !defined(DIST_JM) && !defined(DIST_STAT)
# define DIST_STAT
#endif

using namespace lemon;
using namespace std;
using namespace sds;

typedef ListGraph Graph;
GRAPH_TYPEDEFS(Graph);

typedef Graph::NodeMap<int> IntMap;
typedef Graph::NodeMap<dim2::Point<int> > PosMap;

template <typename _Value>
class PointMap {
public:
  typedef dim2::Point<int> Key;
  typedef _Value Value;

  PointMap(int _width, int _height) 
    : width(_width), height(_height), data(_width * _height) {}

  PointMap(int _width, int _height, int _value) 
    : width(_width), height(_height), data(_width * _height, _value) {}

  Value& operator[](const Key& key) {
    return data[key.x + key.y * width];
  }

  const Value& operator[](const Key& key) const {
    return data[key.x + key.y * width];
  }

  void set(const Key& key, const Value& value) {
    operator[](key) = value;
  }

  int getWidth() const { return width; }
  int getHeight() const { return height; }

private:
  int width, height;
  vector<Value> data;
};

class SimpleCluster {
  friend SimpleCluster operator+(const SimpleCluster&, const SimpleCluster&);
public:
  SimpleCluster() {
    _size = 0;
    for (int k = 0; k < CH; ++k) {
      _sum[k] = 0;
    }
  }

  SimpleCluster& operator=(const Point& p) {
    for (int k = 0; k < CH; ++k) {
      _sum[k] = p[k];
    }
    _size = 1;
    return *this;
  }

  Point center() const {
    Point p;
    for (int k = 0; k < CH; ++k) {
      p[k] = _sum[k] / _size;
    }
    return p;
  }

private:
  int _size;
  double _sum[CH];  
};


SimpleCluster operator+(const SimpleCluster& p, const SimpleCluster& q) {
  SimpleCluster r;
  for (int k = 0; k < CH; ++k) {
    r._sum[k] = p._sum[k] + q._sum[k];
  }
  r._size = p._size + q._size;
  return r;
}

class MCluster {
  friend MCluster operator+(const MCluster&, const MCluster&);
  friend double mdistance(const MCluster&, const MCluster&);
public:
  MCluster() : cluster(0) {}

  MCluster(const MCluster& p) {
    if (p.cluster.get() != 0) {
      cluster.reset(new Cluster());
      *cluster = *(p.cluster);
    } else {
      points = p.points;
    }
  }

  MCluster& operator=(const Point& p) {
    cluster.reset(0);
    points.clear();
    points.push_back(p);
    return *this;
  }

  MCluster& operator=(const MCluster& p) {
    if (p.cluster.get() != 0) {
      cluster.reset(new Cluster());
      *cluster = *(p.cluster);
    } else {
      points = p.points;
    }
    return *this;
  }

  int size() const {
    if (cluster.get() != 0) {
      return cluster->size();
    } else {
      return points.size();
    }
  }

  
private:
  auto_ptr<Cluster> cluster;
  vector<Point> points;
};

MCluster operator+(const MCluster& p, const MCluster& q) {
  MCluster r;
  if (p.cluster.get() != 0 && q.cluster.get() != 0) {
    r.cluster.reset(new Cluster());
    *(r.cluster) = *(p.cluster) + *(q.cluster);
    return r;
  } else if (p.cluster.get() != 0) {
    r.cluster.reset(new Cluster());
    *(r.cluster) = *(p.cluster);
    for (int i = 0; i < (int)q.points.size(); ++i) {
      (*r.cluster) += q.points[i];
    }
    return r;
  } else if (q.cluster.get() != 0) {
    r.cluster.reset(new Cluster());
    *(r.cluster) = *(q.cluster);
    for (int i = 0; i < (int)p.points.size(); ++i) {
      (*r.cluster) += p.points[i];
    }
    return r;
  } else if ((int)p.points.size() + (int)q.points.size() > CH) {
    r.cluster.reset(new Cluster());
    for (int i = 0; i < (int)p.points.size(); ++i) {
      (*r.cluster) += p.points[i];
    }
    for (int i = 0; i < (int)q.points.size(); ++i) {
      (*r.cluster) += q.points[i];
    }
    return r;
  } else {
    for (int i = 0; i < (int)p.points.size(); ++i) {
      r.points.push_back(p.points[i]);
    }
    for (int i = 0; i < (int)q.points.size(); ++i) {
      r.points.push_back(q.points[i]);
    }
    return r;
  }
}

#if defined(DIST_AVG)
#  define DIST_FUN distance_avg
#elif defined(DIST_DIV)
#  define DIST_FUN distance_div
#elif defined(DIST_MAH)
#  define DIST_FUN distance_mah
#elif defined(DIST_BAT)
#  define DIST_FUN distance_bat
#elif defined(DIST_JM)
#  define DIST_FUN distance_jm
#elif defined(DIST_STAT)
#  define DIST_FUN distance_stat
#endif

double mdistance(const MCluster& p, const MCluster& q) {
  if (p.cluster.get() != 0 && q.cluster.get() != 0) {
    return DIST_FUN(*(p.cluster), *(q.cluster));
  } else if (p.cluster.get() != 0) {
    Cluster qv;
    for (int i = 0; i < (int)q.points.size(); ++i) {
      qv += q.points[i];
    }
    return DIST_FUN(*(p.cluster), qv);
  } else if (q.cluster.get() != 0) {
    Cluster pv;
    for (int i = 0; i < (int)p.points.size(); ++i) {
      pv += p.points[i];
    }
    return DIST_FUN(*(q.cluster), pv);
  } else {
    Cluster pv;
    for (int i = 0; i < (int)p.points.size(); ++i) {
      pv += p.points[i];
    }
    Cluster qv;
    for (int i = 0; i < (int)q.points.size(); ++i) {
      qv += q.points[i];
    }
    return DIST_FUN(pv, qv);
  }
}


typedef Graph::NodeMap<MCluster> SegmentMap;


class DistanceMap {
public:
  typedef Edge Key;
  typedef double Value;

  DistanceMap(Graph& _graph, const SegmentMap& _map)
    : graph(_graph), map(_map) {}

  Value operator[](const Key& edge) const {
    Node source = graph.u(edge);
    Node target = graph.v(edge);
    
    return mdistance(map[source], map[target]);
  }

private:
  Graph& graph;

  const SegmentMap& map;
};

void Read(const string& filename, Graph& graph, PosMap& pos, 
	  SegmentMap& seg, vector<SInteger>& size) {
  SFile file(filename);

  SDataView<unsigned char> source[CH];
  for (int k = 0; k < CH; ++k) {
    SDataset dataset = file.getDataset(k);
    source[k] = dataset.getDataView<unsigned char>();
  }
  size = source[0].getSize();

  SDataView<Node> nodes(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Node node = graph.addNode();
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      seg[node] = point;
      pos[node].y = i;
      pos[node].x = j;
      nodes[i][j] = node;
    }
  }

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j + 1 < size[1]; ++j) {
      graph.addEdge(nodes[i][j], nodes[i][j + 1]);
    }
  }
  for (int i = 0; i + 1 < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      graph.addEdge(nodes[i][j], nodes[i + 1][j]);
    }
  }
}

void Write(const string& filename, Graph& graph, PosMap& pos, SegmentMap& seg,
	   UnionFind<PointMap<int> >& uf, vector<SInteger>& size) {
  SFile file(filename, SFile::CREATE);

  IntMap color(graph, -2);
  
  IntMap heapMap(graph, -1);
  BinHeap<int, IntMap> heap(heapMap);
  
  map<int, int> segment;
  
  for (NodeIt it(graph); it != INVALID; ++it) {
    if (seg[it].size() > 1) {
      heap.push(it, countOutArcs(graph, it));
    } else {
      color[it] = 0;
      segment[uf.find(pos[it])] = 0; 
    }
  }

  vector<Node> order;

  while (!heap.empty()) {
    Node node = heap.top();
    heap.pop();
    color[node] = -1;
    order.push_back(node);
    for (IncEdgeIt it(graph, node); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] == -2) {
	heap.decrease(target, heap[target] - 1);
      }
    }
  }

  for (int i = order.size() - 1; i >= 0; --i) {
    set<int> colors;
    for (IncEdgeIt it(graph, order[i]); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] != -1) {
	colors.insert(color[target]);
      }
    }
    for (int j = 1; j <= 6; ++j) {
      if (colors.find(j) == colors.end()) {
	color[order[i]] = j;
	segment[uf.find(pos[order[i]])] = j; 
	break;
      }
    }
  }

  SDataView<char> dataview(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      dataview[i][j] = segment[uf.find(dim2::Point<int>(j, i))];
    }
  }

  SDataset dataset = file.createDataset("segmented", UINT8, size);
  
  dataset.setDataView<char>(dataview);
}


int main(int argc, const char *argv[]) {
  assert(argc == 3);

  Graph graph;

  SegmentMap seg(graph);
  PosMap pos(graph);
  vector<SInteger> size;

  Read(argv[1], graph, pos, seg, size);

  PointMap<int> unionfindMap(size[1], size[0], -1);
  UnionFind<PointMap<int> > unionfind(unionfindMap);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      unionfind.insert(dim2::Point<int>(j, i));
    }
  }

  Graph::EdgeMap<int> heapMap(graph, -1);
  BinHeap<double, Graph::EdgeMap<int> > heap(heapMap);

  DistanceMap distance(graph, seg);
  
  for (EdgeIt it(graph); it != INVALID; ++it) {
    heap.push(it, distance[it]);
  }

  const int STEP =  size[0] * size[1] - 12000;

  for (int i = 0; i < STEP; ++i) {
    if ( (i - 1) * 100 / STEP < i * 100 / STEP) {
      cout << i * 100 / STEP << '%' << endl;
    }

    Edge edge = heap.top();

    Node source = graph.u(edge);
    Node target = graph.v(edge);
    
    Node node = graph.addNode();

    pos[node] = pos[source];

    seg[node] = seg[source] + seg[target];

    unionfind.join(pos[source], pos[target]);

    heap.erase(edge);
    graph.erase(edge);

    set<Node> targets;
    for (IncEdgeIt it(graph, source); it != INVALID; ++it) {
      targets.insert(graph.runningNode(it));
      heap.erase(it);
    }
    graph.erase(source);

    for (IncEdgeIt it(graph, target); it != INVALID; ++it) {
      targets.insert(graph.runningNode(it));
      heap.erase(it);
    }
    graph.erase(target);

    for (set<Node>::iterator it = targets.begin(); it != targets.end(); ++it) {
      Edge edge = graph.addEdge(node, *it);
      heap.push(edge, distance[edge]);
    }
  }

  Write(argv[2], graph, pos, seg, unionfind, size);

  return 0;
}
